import numpy as np
import matplotlib.pyplot as plt


def exp_i(n):
    return np.complex(np.cos(n), np.sin(n))


def calc_omega_powers(n):
    omega = exp_i(-1 * np.pi / n)
    omega_powers = np.power(omega, np.arange(n))
    return omega_powers


def fft_vec(xs, n, start=0, stride=1):
    """cooley-turkey fft"""
    if n == 1:
        return [xs[start]]
    hn, sd = n // 2, stride * 2
    rs = fft_vec(xs, hn, start, sd) + fft_vec(xs, hn, start + stride, sd)
    for i in range(hn):
        e = exp_i(-2 * np.pi * i / n)
        rs[i], rs[i + hn] = rs[i] + e * rs[i + hn], rs[i] - e * rs[i + hn]
        pass
    return rs


def fft_matrix(xs, n):
    rs = fft_matrix_(xs, n, start=0, stride=1)
    return np.array(rs).transpose()


def fft_matrix_(xs, n, start=0, stride=1):
    if n == 1:
        return [xs[:, start]]
    hn, sd = n // 2, stride * 2
    rs = fft_matrix_(xs, hn, start, sd) + fft_matrix_(xs, hn, start + stride, sd)
    for i in range(hn):
        e = exp_i(-2 * np.pi * i / n)
        rs[i], rs[i + hn] = rs[i] + e * rs[i + hn], rs[i] - e * rs[i + hn]
    return rs


def test_fft_vec():
    n = 1024
    vec = np.sin(range(n))
    vec2 = np.sin(np.pi * np.arange(n))
    matrix = np.zeros((2, n))
    matrix[0, :] = vec
    matrix[1, :] = vec2

    rs = fft_matrix(matrix, n)
    # print(rs.shape)
    plt.subplot(4, 1, 1)
    plt.plot(np.arange(n), np.abs(rs[0]))
    plt.subplot(4, 1, 2)
    plt.plot(np.arange(n), np.abs(rs[1]))

    plt.subplot(4, 1, 3)
    result2 = np.fft.fft(matrix)
    plt.plot(np.arange(n), np.abs(result2[0]))
    plt.subplot(4, 1, 4)
    plt.plot(np.arange(n), np.abs(result2[1]))
    plt.show()


if __name__ == '__main__':
    test_fft_vec()
